import numpy as np

from Ashare import *
import matplotlib.pyplot as plt
import efinance as ef

def kechuang(days, useFile, firstCode, secondCode, needBackTest=True, needDrawPic=True):
    disList = []
    dateList = []
    dfSecond = None
    dfFirst = None
    returnDiffDay = 40
    current_date = datetime.datetime.now().date()
    days_ago = current_date - datetime.timedelta(days=days)
    startDate = days_ago.strftime("%Y%m%d")
    endDate = datetime.datetime.now().strftime("%Y%m%d")
    # 科创50
    if useFile:
        dfSecond = pd.read_csv(f'{secondCode}.csv', index_col=0, parse_dates=True)
        dfFirst = pd.read_csv(f'{firstCode}.csv', index_col=0, parse_dates=True)
        print("from file")
    else:
        print("from net")
        #dfSecond = get_price(secondCode, frequency='1d', count=days)
        dfSecond = downloadData(secondCode, startDate, endDate)
        #dfSecond.to_csv(f'{secondCode}.csv', index=True)
        #dfFirst = get_price(firstCode, frequency='1d', count=days)
        dfFirst = downloadData(firstCode, startDate, endDate)
        #dfFirst.to_csv(f'{firstCode}.csv', index=True)

    closeListSecond = dfSecond['close_price']
    numSecond = closeListSecond.size
    # 科创200
    closeListFirst = dfFirst['close_price']
    numFirst = closeListFirst.size
    diffSize = min(numSecond, numFirst)
    print(f'diff size {diffSize}')
    if diffSize == 0:
        return
    closeListSecond = closeListSecond.tail(diffSize)
    closeListFirst = closeListFirst.tail(diffSize)
    for i in range(0, diffSize - returnDiffDay + 1):
        dayBefore40 = closeListSecond[i]
        currentDay = closeListSecond[i + returnDiffDay - 1]
        yieldValueKC50 = (currentDay - dayBefore40) / dayBefore40
        dayBefore40 = closeListFirst[i]
        currentDay = closeListFirst[i + returnDiffDay - 1]
        yieldValueKC200 = (currentDay - dayBefore40) / dayBefore40
        disList.append(round((yieldValueKC200 - yieldValueKC50) * 100, 2))
        dateList.append(closeListSecond.index[i + returnDiffDay - 1].date())

    # 绘制折线图
    if needDrawPic:
        drawPic(firstCode, secondCode, disList, dateList)
    # backTest(8, -12, disList,dateList, closeListFirst, closeListSecond,returnDiffDay)
    if needBackTest:
        for i in range(1, int(max(disList)) + 1):
            for j in range(int(min(disList)), 0):
                backTest(i, j, disList, dateList, closeListFirst, closeListSecond, returnDiffDay)

        print(f'gR({gR}) gTop({gTop}) gBottom({gBottom})')
        for i in range(0, len(gStep)):
            print(gStep[i])
    if needDrawPic:
        drawMoneyPic(firstCode, secondCode, gMoney, dateList)
        print(f"max down({calculate_max_drawdown(gMoney)}) firstMaxDown({calculate_max_drawdown(closeListFirst)}) firstMaxDown({calculate_max_drawdown(closeListSecond)})")
    firstPriceList = closeListFirst
    secondPriceList = closeListSecond
    onlyFirst = (firstPriceList[-1] - firstPriceList[returnDiffDay - 1]) / firstPriceList[returnDiffDay - 1]
    onlySecond = (secondPriceList[-1] - secondPriceList[returnDiffDay - 1]) / secondPriceList[returnDiffDay - 1]
    print(
        f'first start date({firstPriceList.index[returnDiffDay]}) p({firstPriceList[returnDiffDay]}) end{firstPriceList.index[-1]} p({firstPriceList[-1]})')
    print(
        f'second start date({secondPriceList.index[returnDiffDay]}) p({secondPriceList[returnDiffDay]}) end{secondPriceList.index[-1]} p({secondPriceList[-1]})')
    print(f'return onlyFirst({onlyFirst}) onlySecond({onlySecond})\n')
    print(f'today {firstCode}-{secondCode} diff({disList[-1]}) count({len(disList)})')


gTop = 0
gBottom = 0
gR = 0.0
gStep = []
gMoney = []

def downloadData(code,startDate, endDate):
    df = ef.stock.get_quote_history(stock_codes=code, beg=startDate, end=endDate)
    df.rename(columns={'股票代码': 'code','股票名称':'name', '日期': 'date', '收盘': 'close_price'}, inplace=True)
    dfData = df[['code', 'date', 'close_price']]
    dfData.loc[:,'date'] = pd.to_datetime(dfData['date'])
    dfData = dfData.set_index('date')
    dfData.sort_index(inplace=True)
    dfData = dfData.fillna(method='ffill')
    dfData.to_csv(f'rawData/{code}.csv')
    return dfData


def backTest(diffTop, diffBottom, diffList, dateList, firstPriceList, secondPriceList, returnDiffDay):
    initMoney = 10000
    countFirst = initMoney / 2 / firstPriceList[returnDiffDay - 1]
    countSecond = initMoney / 2 / secondPriceList[returnDiffDay - 1]
    moneyList = []
    moneyList.append(initMoney)
    step = []
    maxMoney = initMoney
    minMoney = initMoney
    for i in range(0, len(diffList) - 1):
        if diffList[i] >= diffTop and countFirst != 0:
            money = countFirst * firstPriceList[returnDiffDay + i]
            addSecond = money / secondPriceList[returnDiffDay + i]
            countSecond += addSecond
            countFirst = 0
            step.append(f"{dateList[i]} sell first ")
            # print("sell first \n")
        if diffList[i] <= diffBottom and countSecond != 0:
            money = countSecond * secondPriceList[returnDiffDay + i]
            addFirst = money / firstPriceList[returnDiffDay + i]
            countFirst += addFirst
            countSecond = 0
            step.append(f"{dateList[i]} sell second ")
            # print("sell second\n")
        currentMoney = countFirst * firstPriceList[returnDiffDay + i] + countSecond * secondPriceList[returnDiffDay + i]
        moneyList.append(currentMoney)
        if currentMoney > maxMoney:
            maxMoney = currentMoney
        if currentMoney < minMoney:
            minMoney = currentMoney

    money = countFirst * firstPriceList[-1] + countSecond * secondPriceList[-1]
    r = (money - initMoney) / initMoney
    global gR
    global gTop
    global gBottom
    global gStep
    global gMoney
    if r > gR:
        gR = r
        gTop = diffTop
        gBottom = diffBottom
        gStep = step
        gMoney = moneyList
    # print(f'return r({r})\n')


def drawPic(firstCode, secondCode, dataY, dataX):
    time_series = pd.to_datetime(dataX)
    plt.plot(time_series, dataY)
    # 添加标题和标签
    plt.title(f'{firstCode}-{secondCode} 40-day return difference')
    plt.xlabel('time')
    plt.ylabel('diff')
    # 设置 y 轴刻度更密集
    # plt.yticks(range(min(dataY), max(dataY), 1))  # 设置 y 轴刻度间隔为 5
    # 自动旋转 x 轴标签
    plt.xticks(rotation=90)
    # 显示网格线
    plt.grid(True)
    plt.savefig(f'image/{firstCode}_{secondCode}.jpg', dpi=400)

    # 显示图形
    plt.show()


def drawMoneyPic(firstCode, secondCode, moneyList, dateList):
    time_series = pd.to_datetime(dateList)
    plt.plot(time_series, moneyList)
    # 添加标题和标签
    plt.title(f'{firstCode}-{secondCode} money')
    plt.xlabel('time')
    plt.ylabel('money')
    # 设置 y 轴刻度更密集
    # plt.yticks(range(min(dataY), max(dataY), 1))  # 设置 y 轴刻度间隔为 5
    # 自动旋转 x 轴标签
    plt.xticks(rotation=90)
    # 显示网格线
    plt.grid(True)
    plt.savefig(f'image/{firstCode}_{secondCode}_money.jpg', dpi=400)

    # 显示图形
    plt.show()


def calculate_max_drawdown(prices):
    """
    计算股票的最大回撤。

    参数:
        prices (list 或 numpy array): 股票价格序列。

    返回:
        max_drawdown (float): 最大回撤值，表示为百分比。
        max_drawdown_start (int): 最大回撤开始的索引。
        max_drawdown_end (int): 最大回撤结束的索引。
    """
    # 转换为 numpy 数组以便计算
    # prices = np.array(prices)

    # 计算累计最大值
    cumulative_max = np.maximum.accumulate(prices)

    # 计算回撤值
    drawdowns = (prices - cumulative_max) / cumulative_max

    # 找到最大回撤
    max_drawdown = drawdowns.min()
    # max_drawdown_end = drawdowns.argmin()
    # max_drawdown_start = prices[:max_drawdown_end + 1]

    return max_drawdown


def test():
    df = ef.fund.get_quote_history(fund_code='000699',pz=5)
    #df = ef.fund.get_quote_history(fund_code='000699', beg='20250101', end='20250117')
    print(df)

def dayReport():
    # 科创 200-50 38% [7~-13]
    print("\n科创 200-50 38% [7~-13]\n")
    kechuang(90, False, 'sh000699', 'sh000688', needBackTest=False, needDrawPic=False)
    # 科创 50-白酒 210% [29~-36] 4年
    print("\n科创 50-白酒 210% [29~-36]\n")
    kechuang(90, False, 'sh000688', 'sz399997', needBackTest=False, needDrawPic=False)
    # 信息技术-白酒 615% [31~-36] 7年  1095% [31~-36] 9年
    print("\n信息技术-白酒 615% [31~-36]\n")
    kechuang(90, False, 'sh000993', 'sz399997',needBackTest=False, needDrawPic=False)


if __name__ == '__main__':
    #dayReport()
    # 科创 200-50 38% [7~-13]
    kechuang(40, False, '000699', '000688', needDrawPic=True)
    # 沪深300-中证 500 19%
    #kechuang(200, False,'000300','000905')
    # 沪深300-创业板  91%  [9,-15] [2017,2024]
    # kechuang(4000, False,'sh000300','sz399006')
    # 沪深300-中证1000 39% [17~-17]
    # kechuang(1211, False, 'sh000300', 'sh000852')
    # 信息技术-红利 163% [15~-21]
    # kechuang(2000, False, 'sh000993', 'sh000015')
    # 信息技术-白酒 615% [31~-36] 7年  1095% [31~-36] 9年
    # kechuang(1213, False, 'sh000993', 'sz399997')
    # 科创 50-白酒 210% [29~-36] 4年
    # kechuang(1213, False, 'sh000688', 'sz399997',needDrawPic=True)
    # 科创 50-红利 254% [13~-24] 4年
    # kechuang(1213, False, 'sh000688', 'sh000015')
    # kechuang(1213, False, 'sh000688', 'sh000015')
    # df = get_price('sz399997', frequency='1d', count=5)
    # print(df)
    #test()
